import org.antlr.v4.runtime.*;
import org.antlr.v4.runtime.misc.MultiMap;
import org.antlr.v4.runtime.misc.OrderedHashSet;
import org.antlr.v4.runtime.misc.Pair;
import org.antlr.v4.runtime.tree.ParseTreeWalker;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.List;
import java.util.Set;


class FunctionListener extends CymbolBaseListener{
    Graph graph = new Graph();
    String currentFunctionName = null;

    public void printGraph(){
        MultiMap<String,String> mMap = graph.edges;
        List<Pair<String, String>> pairs = mMap.getPairs();
        System.out.println("=====================");
        for(Pair<String, String> pair : pairs){
            System.out.println(pair.a + "->" + pair.b);
        }
    }

    public String toDot()
    {

        StringBuilder buf = new StringBuilder();
        buf.append("digraph G {\n");
        buf.append(" ranksep=.25;\n");
        buf.append(" edge [arrowsize=.5]\n");
        buf.append(" node [shape=circle, fontname=\"ArialNarrow\",\n");
        buf.append(" fontsize=12, fixedsize=true, height=.45];\n");
        buf.append(" ");
        for (String node : graph.nodes) { // print all nodes first
            buf.append(node);
            buf.append("; ");
        }
        buf.append("\n");
        for (String src : graph.edges.keySet()) {
            for (String trg : graph.edges.get(src)) {
                buf.append(" ");
                buf.append(src);
                buf.append(" -> ");
                buf.append(trg);
                buf.append(";\n");
            }
        }
        buf.append("}\n");
        return buf.toString();
    }

    @Override
    public void enterFunctionDecl(CymbolParser.FunctionDeclContext ctx) {
        currentFunctionName = ctx.ID().getText();
        graph.nodes.add(currentFunctionName);
    }

    @Override
    public void exitCall(CymbolParser.CallContext ctx) {
        String funcName = ctx.ID().getText();
        graph.edge(currentFunctionName,funcName);
    }

    private static class Graph{
        Set<String> nodes = new OrderedHashSet<String>();
        MultiMap<String,String> edges =
                new MultiMap<String,String>();
        public void edge(String source,String target){
            edges.map(source, target);
        }
    }
}

class CallRelation {
    public static void main(String[] args) throws IOException
    {
        InputStream cymbol = new FileInputStream(new File("C:\\usr\\projects\\go_parser\\gen\\test.txt"));
        ANTLRInputStream aInputStream = new ANTLRInputStream(cymbol);
        CymbolLexer lexer = new CymbolLexer(aInputStream);
        CommonTokenStream cts = new CommonTokenStream(lexer);
        CymbolParser parser = new CymbolParser(cts);
        ParseTreeWalker walker = new ParseTreeWalker();
        FunctionListener collector = new FunctionListener();
        walker.walk(collector, parser.file());
        collector.toDot();
    }
}

